package org.hepeng.workx.spring.cloud.netflix.ribbon.loadbalancer;

import com.netflix.loadbalancer.DynamicServerListLoadBalancer;
import com.netflix.loadbalancer.ILoadBalancer;
import com.netflix.loadbalancer.Server;

import java.util.List;

/**
 * @author he peng
 */
public class ThreadIsolationLoadBalancer<T extends Server> extends DynamicServerListLoadBalancer<T> {

    private final ThreadLocal<List<Server>> allServerListThreadLocal = ThreadLocal.withInitial(
            () -> ThreadIsolationLoadBalancer.this.lb.getAllServers());

    private ILoadBalancer lb;

    public ThreadIsolationLoadBalancer(ILoadBalancer lb) {
        this.lb = lb;
    }

    public void setAllServerList(List<Server> serverList) {
        allServerListThreadLocal.set(serverList);
    }

    @Override
    public List<Server> getAllServers() {
        return allServerListThreadLocal.get();
    }

    public void clear() {
        allServerListThreadLocal.remove();
    }
}
